package utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"io"
)

/*

这段代码是一个开箱即用的 AES 加密/解密工具包，支持两种主流模式：
1. 
CBC 模式（传统，需要手动填充、需要 IV）
2. 
GCM 模式（现代，自带认证、无需填充、推荐使用）



模式 		安全级			是否认证  	需要填充 		IV/Nonce 处理         			适用场景

CBC         一般          ❌ 无认证    ✅ PKCS7  		手动生成 IV 并拼在密文前     		兼容老系统 
GCM 		高 			✅ 自带认证 		❌ 无需填充 		自动生成 nonce 并拼在密文前 	新项目首选 

1. CBC 模式

cipherText, _ := aes.EncryptCBC([]byte("hello"))
plainBack, _  := aes.DecryptCBC(cipherText)

加密返回 base64 字符串；解密返回 原始字节。
内部已帮你做 PKCS7 填充 / 去填充。
IV 随机生成并直接拼在密文最前面，解密时自己再拆出来。

2. GCM 模式（推荐）

cipherText, _ := aes.EncryptGCM([]byte("hello"))
plainBack, _  := aes.DecryptGCM(cipherText)

同样返回 base64 字符串。
自带认证标签（tag），篡改数据会解密失败并返回  error 。
nonce 长度 12 字节，随机生成后拼在密文前，解密时自动拆分。



填充函数 pkcs7Pad / pkcs7Unpad

仅 CBC 需要；GCM 不用。
规则：缺 N 个字节就补 N 个  0xN 。
解密后最后一个字节告诉你该删几字节。

六、使用注意

1. 
密钥一定固定长度（16/24/32）。
2. 
GCM 解密失败 = 数据被篡改或密钥错误，不要忽略 error。
3. 
CBC 无认证，无法发现篡改；新项目优先 GCM。
4. 
如果存数据库，直接把返回的 base64 字符串当普通字符串字段即可。


七、
 EncryptGCM / DecryptGCM  是现代加密+防篡改一条龙，优先用；
 EncryptCBC / DecryptCBC  是给老系统向下兼容准备的，记得做好完整性校验（如再加 HMAC）

*/



//把合法长度的 AES 密钥收进来，16/24/32 字节分别对应 AES-128/192/256。
type AESEncryptor struct {
	Key []byte
}

func NewAESEncryptor(key string) (*AESEncryptor, error) {
	// AES密钥长度必须是16, 24或32字节(对应AES-128, AES-192, AES-256)
	keyBytes := []byte(key) //  就是把用户给的字符串密钥转成二进制字节数组
	validLengths := map[int]bool{16: true, 24: true, 32: true}
	if !validLengths[len(keyBytes)] {
		return nil, errors.New("密钥长度必须是16, 24或32字节")
	}
	return &AESEncryptor{Key: keyBytes}, nil
}
func (a *AESEncryptor) EncryptCBC(plaintext []byte) (string, error) {
	block, err := aes.NewCipher(a.Key)
	if err != nil {
		return "", err
	}

	// 填充明文以满足块大小   内部已帮你做 PKCS7 填充 / 去填充
	plaintext = pkcs7Pad(plaintext, aes.BlockSize)

	// 创建存储密文的缓冲区，需要额外空间存储IV  IV 随机生成并直接拼在密文最前面，解密时自己再拆出来。
	ciphertext := make([]byte, aes.BlockSize+len(plaintext))

	// 生成随机初始化向量(IV)
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", err
	}

	// 创建CBC加密器
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext[aes.BlockSize:], plaintext)

	// 返回Base64编码的密文
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// DecryptCBC 使用CBC模式解密
func (a *AESEncryptor) DecryptCBC(ciphertext string) ([]byte, error) {
	// 解码Base64密文
	decoded, err := base64.StdEncoding.DecodeString(ciphertext)
	if err != nil {
		return nil, err
	}

	block, err := aes.NewCipher(a.Key)
	if err != nil {
		return nil, err
	}

	if len(decoded) < aes.BlockSize {
		return nil, errors.New("密文太短")
	}

	// 提取IV和实际密文
	iv := decoded[:aes.BlockSize]
	decoded = decoded[aes.BlockSize:]

	// 创建CBC解密器
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(decoded, decoded)

	// 去除填充
	return pkcs7Unpad(decoded), nil
}

// EncryptGCM 使用GCM模式加密（推荐，更安全）
func (a *AESEncryptor) EncryptGCM(plaintext []byte) (string, error) {
	block, err := aes.NewCipher(a.Key)
	if err != nil {
		return "", err
	}

	// 创建GCM模式
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	// 生成随机nonce
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	// 加密并认证
	ciphertext := gcm.Seal(nonce, nonce, plaintext, nil)

	// 返回Base64编码的密文
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// DecryptGCM 使用GCM模式解密
func (a *AESEncryptor) DecryptGCM(ciphertext string) ([]byte, error) {
	// 解码Base64密文
	decoded, err := base64.StdEncoding.DecodeString(ciphertext)
	if err != nil {
		return nil, err
	}

	block, err := aes.NewCipher(a.Key)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonceSize := gcm.NonceSize()
	if len(decoded) < nonceSize {
		return nil, errors.New("密文太短")
	}

	// 提取nonce和实际密文
	nonce, ciphertextBytes := decoded[:nonceSize], decoded[nonceSize:]

	// 解密并验证
	return gcm.Open(nil, nonce, ciphertextBytes, nil)
}

// pkcs7Pad PKCS7填充
func pkcs7Pad(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// pkcs7Unpad PKCS7去除填充
func pkcs7Unpad(data []byte) []byte {
	length := len(data)
	unpadding := int(data[length-1])
	return data[:(length - unpadding)]
}
